import { Client, IMessage, StompSubscription } from '@stomp/stompjs'
import SockJS from 'sockjs-client'

export type MessageHandler = (message: any) => void

class WebSocketManager {
  private client: Client | null = null
  private subscriptions: Map<string, StompSubscription> = new Map()
  private messageHandlers: Map<string, MessageHandler[]> = new Map()
  private reconnectAttempts = 0
  private maxReconnectAttempts = 5
  private reconnectDelay = 3000

  /**
   * 连接WebSocket
   */
  connect(token: string): Promise<void> {
    return new Promise((resolve, reject) => {
      this.client = new Client({
        webSocketFactory: () => {
          return new SockJS(`${import.meta.env.VITE_WS_URL || 'http://localhost:8080'}/ws`)
        },
        
        connectHeaders: {
          Authorization: `Bearer ${token}`
        },

        debug: (str) => {
          if (import.meta.env.DEV) {
            console.log('[STOMP]', str)
          }
        },

        reconnectDelay: this.reconnectDelay,
        heartbeatIncoming: 4000,
        heartbeatOutgoing: 4000,

        onConnect: () => {
          console.log('WebSocket connected')
          this.reconnectAttempts = 0
          
          // 重新订阅所有主题
          this.resubscribeAll()
          
          resolve()
        },

        onStompError: (frame) => {
          console.error('STOMP error', frame)
          reject(new Error(frame.headers['message'] || 'STOMP error'))
        },

        onWebSocketError: (error) => {
          console.error('WebSocket error', error)
        },

        onDisconnect: () => {
          console.log('WebSocket disconnected')
          this.attemptReconnect()
        }
      })

      this.client.activate()
    })
  }

  /**
   * 断开连接
   */
  disconnect() {
    if (this.client) {
      this.subscriptions.clear()
      this.messageHandlers.clear()
      this.client.deactivate()
      this.client = null
    }
  }

  /**
   * 订阅主题
   */
  subscribe(destination: string, handler: MessageHandler) {
    if (!this.client || !this.client.connected) {
      console.warn('WebSocket not connected, saving handler for later')
      this.addMessageHandler(destination, handler)
      return
    }

    // 如果已经订阅过，直接添加处理器
    if (this.subscriptions.has(destination)) {
      this.addMessageHandler(destination, handler)
      return
    }

    const subscription = this.client.subscribe(destination, (message: IMessage) => {
      try {
        const data = JSON.parse(message.body)
        this.handleMessage(destination, data)
      } catch (error) {
        console.error('Failed to parse message', error)
      }
    })

    this.subscriptions.set(destination, subscription)
    this.addMessageHandler(destination, handler)
  }

  /**
   * 取消订阅
   */
  unsubscribe(destination: string, handler?: MessageHandler) {
    if (handler) {
      // 移除特定处理器
      const handlers = this.messageHandlers.get(destination) || []
      const index = handlers.indexOf(handler)
      if (index > -1) {
        handlers.splice(index, 1)
      }

      // 如果没有处理器了，取消订阅
      if (handlers.length === 0) {
        const subscription = this.subscriptions.get(destination)
        if (subscription) {
          subscription.unsubscribe()
          this.subscriptions.delete(destination)
        }
        this.messageHandlers.delete(destination)
      }
    } else {
      // 移除所有处理器并取消订阅
      const subscription = this.subscriptions.get(destination)
      if (subscription) {
        subscription.unsubscribe()
        this.subscriptions.delete(destination)
      }
      this.messageHandlers.delete(destination)
    }
  }

  /**
   * 发送消息
   */
  send(destination: string, body: any) {
    if (!this.client || !this.client.connected) {
      console.error('WebSocket not connected')
      return
    }

    this.client.publish({
      destination,
      body: JSON.stringify(body)
    })
  }

  /**
   * 订阅会话消息
   */
  subscribeToConversation(conversationId: string, handler: MessageHandler) {
    this.subscribe(`/topic/conversation/${conversationId}`, handler)
  }

  /**
   * 取消订阅会话消息
   */
  unsubscribeFromConversation(conversationId: string, handler?: MessageHandler) {
    this.unsubscribe(`/topic/conversation/${conversationId}`, handler)
  }

  /**
   * 订阅用户私有消息
   */
  subscribeToUser(userId: string, channel: string, handler: MessageHandler) {
    this.subscribe(`/user/${userId}/queue/${channel}`, handler)
  }

  /**
   * 发送正在输入通知
   */
  sendTyping(conversationId: string, isTyping: boolean) {
    this.send('/app/typing', {
      conversationId,
      isTyping
    })
  }

  /**
   * 更新在线状态
   */
  updateStatus(status: string) {
    this.send('/app/status.update', status)
  }

  // 私有方法

  private addMessageHandler(destination: string, handler: MessageHandler) {
    const handlers = this.messageHandlers.get(destination) || []
    handlers.push(handler)
    this.messageHandlers.set(destination, handlers)
  }

  private handleMessage(destination: string, data: any) {
    const handlers = this.messageHandlers.get(destination) || []
    handlers.forEach(handler => {
      try {
        handler(data)
      } catch (error) {
        console.error('Message handler error', error)
      }
    })
  }

  private resubscribeAll() {
    const destinations = Array.from(this.messageHandlers.keys())
    this.subscriptions.clear()

    destinations.forEach(destination => {
      const handlers = this.messageHandlers.get(destination) || []
      if (handlers.length > 0 && this.client) {
        const subscription = this.client.subscribe(destination, (message: IMessage) => {
          try {
            const data = JSON.parse(message.body)
            this.handleMessage(destination, data)
          } catch (error) {
            console.error('Failed to parse message', error)
          }
        })
        this.subscriptions.set(destination, subscription)
      }
    })
  }

  private attemptReconnect() {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnect attempts reached')
      return
    }

    this.reconnectAttempts++
    console.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`)
  }
}

export const wsManager = new WebSocketManager()
